package com.sprint.questai.pdfTextExtraction;

import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * PDFTextExtractionServiceImpl 类：实现 PDFTextExtractionService 接口的具体方法。
 */
@Service
public class PDFTextExtractionImpl implements PDFTextExtraction {
    private static final Logger log = LoggerFactory.getLogger(PDFTextExtractionImpl.class);
    private final ITesseract tesseract;
    private final ExecutorService executorService;

    public PDFTextExtractionImpl() {
        tesseract = new Tesseract();
        tesseract.setDatapath("src/main/resources/tessdata");
        tesseract.setLanguage("chi_sim");
        int numberOfCores = Runtime.getRuntime().availableProcessors();
        executorService = Executors.newFixedThreadPool(numberOfCores);
    }

    @Override
    public String extractSelectableText(File file) {
        try (PDDocument document = PDDocument.load(file)) {
            PDFTextStripper pdfStripper = new PDFTextStripper();
            return pdfStripper.getText(document);
        } catch (IOException e) {
            log.error("Failed to extract selectable text from file: {}", file.getName(), e);
            throw new RuntimeException("Text extraction failed", e);
        }
    }

    @Override
    public String extractImageText(File file) {
        try (PDDocument document = PDDocument.load(file)) {
            PDFRenderer pdfRenderer = new PDFRenderer(document);
            StringBuilder imageText = new StringBuilder();

            for (int page = 0; page < document.getNumberOfPages(); page++) {
                BufferedImage bufferedImage = pdfRenderer.renderImageWithDPI(page, 300);
                String ocrText = tesseract.doOCR(bufferedImage);
                imageText.append(ocrText);
            }
            return imageText.toString();
        } catch (IOException | TesseractException e) {
            log.error("Failed to extract image text from file: {}", file.getName(), e);
            throw new RuntimeException("Image text extraction failed", e);
        }
    }

    @Override
    public String extractFullText(File file) {
        String selectableText = extractSelectableText(file);
        String imageText = extractImageText(file);
        return selectableText + "\n" + imageText;
    }

    @Override
    public String extractImageTextM(File file) {
        try (PDDocument document = PDDocument.load(file)) {
            PDFRenderer pdfRenderer = new PDFRenderer(document);
            int numberOfPages = document.getNumberOfPages();
            List<Future<String>> futures = new ArrayList<>();
            int maxThreads = Math.min(numberOfPages, 4);  // 限制最大线程数为4
            ExecutorService executorService = Executors.newFixedThreadPool(maxThreads);

            for (int page = 0; page < numberOfPages; page++) {
                final int currentPage = page;
                futures.add(executorService.submit(() -> {
                    BufferedImage bufferedImage = pdfRenderer.renderImageWithDPI(currentPage, 300);
                    ITesseract tesseract = new Tesseract();
                    tesseract.setDatapath("src/main/resources/tessdata");
                    tesseract.setLanguage("chi_sim");
                    return tesseract.doOCR(bufferedImage);
                }));
            }

            StringBuilder imageText = new StringBuilder();
            for (Future<String> future : futures) {
                imageText.append(future.get());
            }

            executorService.shutdown();
            return imageText.toString();
        } catch (IOException | InterruptedException | ExecutionException e) {
            log.error("Failed to extract image text from file: {}", file.getName(), e);
            throw new RuntimeException("Image text extraction failed", e);
        }
    }


    @Override
    public String extractImageTextWithoutWhitespace(File file) {
        try (PDDocument document = PDDocument.load(file)) {
            PDFRenderer pdfRenderer = new PDFRenderer(document);
            StringBuilder imageText = new StringBuilder();

            for (int page = 0; page < document.getNumberOfPages(); page++) {
                BufferedImage bufferedImage = pdfRenderer.renderImageWithDPI(page, 300);
                String ocrText = tesseract.doOCR(bufferedImage);
                imageText.append(ocrText.replaceAll("\\s+", ""));
            }
            return imageText.toString();
        } catch (IOException | TesseractException e) {
            log.error("Failed to extract image text without whitespace from file: {}", file.getName(), e);
            throw new RuntimeException("Image text extraction without whitespace failed", e);
        }
    }

//    public static void main(String[] args) {
//        PDFTextExtraction textExtractionService = new PDFTextExtractionImpl();
//
//        // 测试文件路径
//        String filePath = "D:\\Project\\QuestAi\\quest-ai-backend\\src\\main\\resources\\legal_data\\PDF\\上海某石油化工有限公司诉山西某科技有限公司买卖合同纠纷案(FBM-CLI.C.546029156).pdf";
//        File file = new File(filePath);
//        if (!file.exists()) {
//            System.err.println("Test file not found: " + filePath);
//            return;
//        }
//
//        try {
//            String imageText = textExtractionService.extractImageTextM(file);
//            System.out.println("Image Text:\n" + imageText);
//        } catch (RuntimeException e) {
//            System.err.println("Failed to process OCR: " + e.getMessage());
//        }
//    }
}
